
import torch
from torch.utils.data import DataLoader
from model import FiberVibrationClassifier
from train import train_model
from dataset import FiberDataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Dataset
train_dataset = FiberDataset(csv_file='train.csv')
val_dataset = FiberDataset(csv_file='val.csv')
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32)

# Model, loss, optimizer
model = FiberVibrationClassifier().to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Train
train_model(model, train_loader, val_loader, num_epochs=30, 
            criterion=criterion, optimizer=optimizer, 
            device=device, save_path='fiber_best_model.pth')
